#!/usr/bin/env python
# -*- coding:utf8 -*-

import os, os.path
import numpy as np
import codecs
import math

#-------------------------------------------------------------------------------------------------------------
# 该部分使用DFT与IDFT的复数极坐标方式来对数组进行分解和合成
def forier_ufunc(nharmonic, Nsample, signs):
    # 输入为谐波倍数，样本数，符号
    # 符号是为了能复用该函数，DFT用-1，IDFT用1
    def forierfunc(ipoint):
        # 输入为要计算采样点的序数
        return np.e**complex(0,(signs * 2*ipoint*nharmonic*np.pi/Nsample))

    return np.frompyfunc(forierfunc, 1, 1)

def DFT1(samples):
    # 使用DFT的复数极坐标的公式来计算和输出结果
    # 输入为一个序列，样本数为偶数
    # 输出为一个序列，存放DFT结果：复数类型
    # 将采样点序列转换为numpy的数组类型
    Nsam = len(samples)
    xn = np.array(samples)
    # 计算e**-j2knPi/N的序列，其中n从0到N-1
    # 先构造一个0到N-1的序列或者数组
    tmparray = np.arange(0,Nsam,1)
    # 然后对每个元素进行带入运算，使用numpy的frompyfunc函数，使之可以处理数组

    # 逐个计算每个频域的值，从0到N-1
    fvalue = []
    for isam in xrange(Nsam):
        # 先计算该样本点对应的频率序列点的值，也就是基函数部分，然后与scale逐个相乘取平均即可
        basicvalues = forier_ufunc(isam, Nsam, -1)(tmparray)
        # 逐个相乘取平均
        fvalue.append(float(1)/Nsam*sum(basicvalues*xn))

    # 返回DFT值，都是复数值
    return fvalue

def IDFT1(samples):
    # 与DFT1相反，根据DFT序列来还原原始信号
    # 输入为复数序列
    # 运算与DFT差不多，signs不一样，而且后来不用取平均
    Nsam = len(samples)
    xn = np.array(samples)
    # 计算e**-j2knPi/N的序列，其中n从0到N-1
    # 先构造一个0到N-1的序列或者数组
    tmparray = np.arange(0,Nsam,1)
    # 然后对每个元素进行带入运算，使用numpy的frompyfunc函数，使之可以处理数组

    # 逐个计算每个频域的值，从0到N-1
    fvalue = []
    for isam in xrange(Nsam):
        # 先计算该样本点对应的频率序列点的值，也就是基函数部分，然后与scale逐个相乘取平均即可
        basicvalues = forier_ufunc(isam, Nsam, 1)(tmparray)
        # 逐个相乘取平均
        fvalue.append(sum(basicvalues*xn))

    # 返回DFT值，都是复数值
    return fvalue

#--------------------------------------------------------------------------------------------------------------
# 该部分使用DFT和IDFT的实数方式来计算
# 需要分别计算实部和虚部；在运用IDFT的时候，还要乘以系数
# 样本数目要是偶数


def cosine_ufunc(ksample, Nsam):
    def cosfunc(iseq):
        return math.cos(2*ksample*iseq*np.pi/Nsam)

    return np.frompyfunc(cosfunc, 1, 1)


def sinoid_ufunc(ksample, Nsam):
    def sinfunc(iseq):
        return math.sin(2*ksample*iseq*np.pi/Nsam)

    return np.frompyfunc(sinfunc, 1, 1)


def DFT2(samples):
    # 返回两个数组，分别是实部和虚部
    # 得到样本数
    Nsam = len(samples)
    # 实部
    realvalues = []
    # 虚部
    imagvalues = []
    tmparray = np.array(xrange(Nsam))
    for kharm in xrange(int(Nsam/2+1)):
        # 计算实部
        realvalues.append(sum(samples*cosine_ufunc(kharm, Nsam)(tmparray)))
        # 计算虚部
        imagvalues.append(-1*sum(samples*sinoid_ufunc(kharm, Nsam)(tmparray)))

    return np.array(realvalues), np.array(imagvalues)


def IDFT2(realvalues, imagvalues):
    # 输入为实部数组和虚部数组
    # 输出为合成信号的采样数组
    # 对实部数组首尾两个数值先除以2
    realvalues[0] = float(realvalues[0])/2
    realvalues[-1] = float(realvalues[-1])/2
    # 对所有实部数组的元素都除以N/2
    Nsam = (len(realvalues)-1)*2
    realvalues = realvalues*2/Nsam
    # 对所有虚部数组的元素都除以-N/2
    imagvalues = imagvalues*(-2)/Nsam
    # 计算信号采样值
    tmparray = np.array(xrange(Nsam/2+1))
    signalvalues = []
    for isample in xrange(Nsam):
        # 先计算正弦和余弦序列
        cosvalues = cosine_ufunc(isample, Nsam)(tmparray)
        sinvalues = sinoid_ufunc(isample, Nsam)(tmparray)
        # 合并起来计算信号采样值
        signalvalues.append(sum(realvalues*cosvalues+imagvalues*sinvalues))

    return signalvalues


#-------------------------------------------------------------------------------------
# 该部分使用DFT的复数的二维坐标表示来运算
def complex_ufunc():
    def toImag(realnum):
        return complex(0, realnum)

    return np.frompyfunc(toImag, 1, 1)


def DFT3(samples):
    # 可以借用cosine_ufunc和sinoid_ufunc
    # 输入是N-point的样本序列，N为偶数
    # 输入是N-piont的DFT结果，复数数组
    # 首先，计算每一项的基函数，也就是余弦和正弦组成的复数
    dftvalues = []
    Nsam = len(samples)
    tmparray = np.array(xrange(Nsam))
    for ksam in xrange(Nsam):
        cosvalues = cosine_ufunc(ksam, Nsam)(tmparray)
        sinvalues = complex_ufunc()(sinoid_ufunc(ksam, Nsam)(tmparray))
        basicvalues = cosvalues - sinvalues
        dftvalues.append(sum(samples*basicvalues)/Nsam)

    return np.array(dftvalues)



def IDFT3(samples):
    # 输入是DFT的结果，也即是复数数组
    # 输出是原始信号，也是复数数组，不过虚部应该都是0
    # 首先，计算每一项的基函数，也就是余弦和正弦组成的复数
    idftvalues = []
    Nsam = len(samples)
    tmparray = np.array(xrange(Nsam))
    for ksam in xrange(Nsam):
        cosvalues = cosine_ufunc(ksam, Nsam)(tmparray)
        sinvalues = sinoid_ufunc(ksam, Nsam)(tmparray)
        basicvalues1 = cosvalues + complex_ufunc()(sinvalues)
        basicvalues2 = sinvalues - complex_ufunc()(cosvalues)
        idftvalues.append(sum(samples.real*basicvalues1-samples.imag*basicvalues2))

    return idftvalues
#------------------------------------------------------------------------------------
# 该部分用于测试
def test():
    test_DFT3andIDFT3()

def test_DFT1andIDFT1():
    testlist = xrange(1024)
    dftlist = DFT1(testlist)
    psave = codecs.open(r"d:\test.txt", "w", encoding="utf8")
    for ivalue in dftlist:
        psave.write("%r\n" % ivalue)
    psave.close()
    pidft = codecs.open(r"d:\idft.txt", "w", encoding="utf8")
    idftlist = IDFT1(dftlist)
    for ivalue in idftlist:
        pidft.write("%r\n" % ivalue)
    pidft.close()


def test_DFT2andIDFT2():
    testlist = xrange(1024)
    realvalues, imagvalues = DFT2(testlist)
    psave = codecs.open(r"d:\dft2.txt", "w", encoding="utf8")
    psave.write("Real Part:\n")
    for ivalue in realvalues:
        psave.write("%r\n" % ivalue)
    psave.write("\n\nImagine Part:\n")
    for ivalue in imagvalues:
        psave.write("%r\n" % ivalue)
    psave.close()
    pidft = codecs.open(r"d:\idft2.txt", "w", encoding="utf8")
    idftlist = IDFT2(realvalues, imagvalues)
    for ivalue in idftlist:
        pidft.write("%r\n" % ivalue)
    pidft.close()


def test_DFT3andIDFT3():
    testlist = xrange(1024)
    dftlist = DFT3(testlist)
    psave = codecs.open(r"d:\dft3.txt", "w", encoding="utf8")
    for ivalue in dftlist:
        psave.write("%r\n" % ivalue)
    psave.close()
    pidft = codecs.open(r"d:\idft3.txt", "w", encoding="utf8")
    idftlist = IDFT3(dftlist)
    for idft in idftlist:
        pidft.write("%r\n" % idft)
    pidft.close()


#-------------------------------------------------------------------------------------------------------------
if __name__ == "__main__":
    test()